import logging
import sys
from collections import defaultdict

from rocket.estimator.hook.hook import AfterIterHook, BeforeRunHook, AfterEpochHook


class LoggerHook(AfterIterHook, BeforeRunHook, AfterEpochHook):
    log_pattern = "\n[%s]: %s"

    def __init__(self, name="LoggerHook", interval=100):
        super().__init__(name)
        self.interval = interval
        self.counter = 0
        self.logger = self._config_logger()

    def before_run(self, container):
        container.put_extra("log", defaultdict(list))

    def after_iter(self, container):
        self.counter += 1
        if self.counter % self.interval == 0:
            self._flush_log(container)

    def after_epoch(self, container):
        self._flush_log(container)

    def _flush_log(self, container):
        log_str = "[Epoch%d][Iter%d]: " % (container.epoch, container.iter)
        for hook_name, logs in container.get_extra("log").items():
            if len(logs) > 0:
                log_str += self.log_pattern % (hook_name, logs[-1])
        self.logger.info(log_str)
        container.put_extra("log", defaultdict(list))

    @staticmethod
    def _config_logger():
        logger = logging.getLogger(__name__)
        logger.setLevel(level=logging.INFO)
        handler = logging.StreamHandler(sys.stdout)
        formatter = logging.Formatter('[%(asctime)s][%(name)s][%(levelname)s]%(message)s')
        handler.setFormatter(formatter)
        logger.addHandler(handler)

        return logger
